# 2024年11月27日-10:03
# 目标，然后仿真阵元数多点，互耦系数多点的情况，心理有个大概的印象，再来实现我们采样方法的互耦抑制
# 目标RMSE
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.linalg import eigh
from scipy.linalg import toeplitz

def printD(tmat):
    print(pd.DataFrame(tmat))
def generate_mutual_coupling_matrix(M,qu):
    """
    生成一个 MxM 的互耦矩阵
    对角线元素为 1，非对角线元素随着距离增加快速衰减，并且是复数
    """
    C = np.eye(M, dtype=complex)  # 初始化为单位矩阵

    for i in range(M):
        for j in range(i + 1, M):
            distance = abs(i - j)
            real_part = np.exp(-distance / 2)  # 实部快速衰减
            imag_part = np.exp(-distance / 2)  # 虚部快速衰减
            C[i, j] = real_part + 1j * imag_part
            C[j, i] = C[i, j]  # 对称性
    
    # 设置左下角和右上角部分为零
    for i in range(M):
        for j in range(M):
            if (j >= M - qu and i <= j - (M - qu)) or (i >= M - qu and j <= i - (M - qu)):
                C[i, j] = 0

    return C

def steering_vector(M, angle, wavelength=1, d=0.5):
    theta = np.deg2rad(angle)
    return np.exp(1j * 2 * np.pi * d * np.arange(M) * np.sin(theta) / wavelength)

def generate_complex_matrix(M):
    """
    生成 MxM 的矩阵 C
    对角线元素为 1，其他非对角线元素为复数，并且所有斜对角线元素相同
    """
    # 初始化矩阵 C 为 MxM 的单位矩阵
    C = np.eye(M, dtype=complex)

    # 生成复数元素
    for i in range(1, M):
        real_part = np.random.rand()  # 随机实部
        imag_part = np.random.rand()  # 随机虚部
        complex_value = real_part + 1j * imag_part  # 生成复数

        for j in range(M - i):
            C[j, j + i] = complex_value  # 上三角斜对角线
            C[j + i, j] = complex_value  # 下三角斜对角线
                
        # if(i==M-1 and j==0) or (i==0 and j==M-1):
        #     C[i,j] = 0
        #     C[j,i] = 0
            
    # for i in range(M):
    #     for j in range(M):
    #         if(j!=0):
    #             C[i,j] = 0
    #         if (i==0 and j==M-2) or (i==1 and j==M-1):
    #             C[i,j] = 0
    #             C[j,i] = 0
                
    return C

def awgn(signal, snr_db):
    snr_linear = 10 ** (snr_db / 10.0)
    power_signal = np.sum(np.abs(signal) ** 2) / signal.size
    power_noise = power_signal / snr_linear
    noise = np.sqrt(power_noise / 2) * (np.random.randn(*signal.shape) + 1j * np.random.randn(*signal.shape))
    return signal + noise

def generate_mat_J(mat_C, q, M):
    mat_J = np.zeros((M, M), dtype=complex)
    for i in range(M):
        for j in range(M):
            if mat_C[i, j] == q:
                mat_J[i, j] = 1
    return mat_J

import pandas as pd
def gMat_T(M, K, DoA_i, mat_C,Q_len):
    d = 0.5  # number of target
    Geo_Tx = np.arange(M)  # geometry of Tx
    rows = mat_C.shape[0]
    Q = Q_len
    Mat_T = np.zeros((M, Q), dtype=complex)
    for q in range(Q):
        mat_J = generate_mat_J(mat_C, q + 1, rows)  # MATLAB 1-indexed, Python 0-indexed
        #df_j = pd.DataFrame(mat_J)
        #print(df_j)
        At = np.exp(1j * 2 * np.pi * d * Geo_Tx * np.sin(np.deg2rad(DoA_i)))
        Jq = mat_J @ At
        Mat_T[:, q] = Jq
    
    return Mat_T

# MC互耦阵列的索引
def gMatC(size):
    mat = np.zeros((size, size), dtype=float)
    last_s = 1
    for i in range(size):
        for j in range(size):
            if i==j:
                mat[i,j]=1
            #elif j < i:
                # mat[i, j] = last_s + 1
                # last_s = mat[i, j]
                # mat[j,i] = last_s
            else:
                mat[i, j] = 1+abs(i - j)
            # if(j!=0 and j!=i):
            #     mat[i,j] = 0
    
    return mat

from scipy.signal import find_peaks
def find_local_maxima(G, theta_list):
    """
    找出数组 G 中的极大值，并返回其对应的下标和 DOA 值，按极大值大小降序排列下标。
    
    参数：
        G (numpy array): 目标函数的值，包含所有搜索网格的响应。
        theta_list (list or numpy array): 对应搜索网格的角度值。
        
    返回：
        peaks_indices (numpy array): 极大值对应的下标，按极大值大小降序排列。
        peaks_DOA (list): 极大值对应的 DOA 值。
    """
    # 找到极大值的下标
    peaks_indices, _ = find_peaks(G)
    
    # 提取相应的极大值
    peaks_values = G[peaks_indices]
    
    # 按照极大值的大小降序排列下标
    sorted_indices = np.argsort(peaks_values)[::-1]
    
    # 根据降序排序后的下标获取对应的极大值和 DOA 值
    sorted_peaks_indices = peaks_indices[sorted_indices]
    sorted_peaks_values = peaks_values[sorted_indices]
    peaks_DOA = [theta_list[i] for i in sorted_peaks_indices]
    
    return sorted_peaks_indices, peaks_DOA

# def find_local_maxima(G, theta_list):
#     """
#     找出数组 G 中的极大值，并返回其对应的下标和 DOA 值。
    
#     参数：
#         G (numpy array): 目标函数的值，包含所有搜索网格的响应。
#         theta_list (list or numpy array): 对应搜索网格的角度值。
        
#     返回：
#         peaks_indices (numpy array): 极大值对应的下标。
#         peaks_DOA (list): 极大值对应的 DOA 值。
#     """
#     # 找到极大值的下标
#     peaks_indices, _ = find_peaks(G)
#     # 提取相应的 DOA 值
#     peaks_DOA = [theta_list[i] for i in peaks_indices]
    
#     return peaks_indices, peaks_DOA

def huou_way(M,K,En,search_grid,MCM,c,q):
    mat_C = gMatC(M)
    Q_len = M-q#-K #M-K
    e = np.eye(Q_len, 1)  # d 约束条件
    G = np.zeros(len(search_grid))  # store the search values
    maxV = -999
    max_DOA = 999
    theta_lst = []
    G2 = np.zeros(len(search_grid))
    for idx,theta in enumerate(search_grid):
        Mat_T = gMat_T(M, K, theta, mat_C,Q_len)
        Q_mat = Mat_T.conj().T @ (En @ En.conj().T) @ Mat_T
        res = np.abs(e.conj().T @ np.linalg.inv(Q_mat) @ e)
        G[idx] = res  # g function
        theta_lst.append(theta)
        if res>maxV:
            maxV = res
            max_DOA = theta
    inds,DOA_lst = find_local_maxima(G,theta_lst)
    #print(DOA_lst)
    c1 = 1
        
        
    return G,c1,max_DOA,DOA_lst
    
def MUSIC_way(M,K,En,search_grid):
    Geo_Tx = np.arange(M)
    G = np.zeros(len(search_grid))
    for idx,theta in enumerate(search_grid):
        At = np.exp(1j * 2 * np.pi * d * Geo_Tx * np.sin(np.deg2rad(theta)))  
        G[idx] = abs(1/(At.conj().T@En@En.conj().T@At))  
    return G

def sort_doa(doa_list):
    sorted_doa_list = sorted(doa_list, reverse=True)
    return sorted_doa_list

def find_top_n_angles(doa_list, n):
    """
    找出 doa_list 中最大的 n 个角度，降序排列。
    
    参数：
        doa_list (list or numpy array): 包含所有 DOA 值的列表。
        n (int): 需要返回的最大角度的数量。
        
    返回：
        top_n_angles (list): 最大的 n 个角度，按降序排列。
    """
    # 对 doa_list 进行降序排序，并返回前 n 个值
    sorted_doa_list = sorted(doa_list, reverse=True)
    top_n_angles = sorted_doa_list[:n]
    
    return top_n_angles

def calculate_rmse(estimated_doa, real_doa):
    """
    计算估计的 DOA 与真实 DOA 之间的 RMSE。
    
    参数：
        estimated_doa (list or numpy array): 估计的 DOA 值列表。
        real_doa (list or numpy array): 真实的 DOA 值列表。
        
    返回：
        rmse (float): 估计 DOA 与真实 DOA 之间的均方根误差。
    """
    # 确保两个列表长度相等
    assert len(estimated_doa) == len(real_doa), "估计的 DOA 和真实 DOA 长度必须相等"
    
    # 计算 RMSE
    squared_errors = [(e - r) ** 2 for e, r in zip(estimated_doa, real_doa)]
    rmse = np.sqrt(np.mean(squared_errors))
    
    return rmse

# def cal_C_again(M,K,mat_C,Q_len,max_DOA,e):
#     Mat_T = gMat_T(M, K, max_DOA, mat_C,Q_len)
#     Q_mat = Mat_T.conj().T @ (En @ En.conj().T) @ Mat_T
#     c1 = (np.linalg.inv(Q_mat) @ e)/(e.conj().T @ np.linalg.inv(Q_mat) @ e)

def Cal_mcm(phi,M,w,mat_C,En,doa_Gu_lst,Q_len,K):
    Q_sum = 0
    for idx,doai in enumerate(doa_Gu_lst):
        Mat_T = gMat_T(M, K,doai, mat_C,Q_len)
        Q_mat = Mat_T.conj().T @ (En @ En.conj().T) @ Mat_T
        Q_sum += Q_mat
    c = (np.linalg.inv(Q_sum) @ w) /(w.conj().T @ Q_sum @ w)
    c = c/c[0]
    cost_val2 = abs(c.conj().T @ Q_sum @ c) # 越小越好
    psi = np.kron(c,np.eye(M))
    C = phi @ psi
    return C

def gMat_J_all(M,mat_C,Q):
    Mat_all_J = np.zeros((M, M * Q))
    rows = mat_C.shape[0]
    for q in range(Q):
        mat_J = generate_mat_J(mat_C, q + 1, rows)
        # 将 mat_J 按列堆叠到 Mat_all_J 中
        start_col = q * M
        end_col = (q + 1) * M
        Mat_all_J[:, start_col:end_col] = mat_J
    return Mat_all_J

def calculate_rmse_complex(estimated_doa, real_doa):
    """
    计算虚数估计 DOA 与真实 DOA 之间的 RMSE。
    """
    # 确保两个列表长度相等
    assert len(estimated_doa) == len(real_doa), "估计的 DOA 和真实 DOA 长度必须相等"
    
    # 分别计算实部和虚部的误差
    real_part_error = [(e.real - r.real) ** 2 for e, r in zip(estimated_doa, real_doa)]
    imag_part_error = [(e.imag - r.imag) ** 2 for e, r in zip(estimated_doa, real_doa)]
    
    # 计算总的 RMSE
    rmse = np.sqrt(np.mean(real_part_error + imag_part_error))
    
    return rmse

def main(SNR_i): 
    # 参数设置
    M = 10  # 天线的数量
    lambda_ = 1  # 波长
    d = 0.5  # 阵元间距
    theta = np.array([25,50,75])  # DOA #[25,50,75]
    K = len(theta)  # 信号个数
    L = 100  # 快拍数
    geo = np.arange(M)  # 与单位阵间隔天线的个数
    A = np.exp(1j * 2 * np.pi * d * np.outer(geo,np.sin(np.deg2rad(theta))))   # 原本的方向矩阵
    S = np.random.randn(K, L) + 1j * np.random.randn(K, L)  # 信号矩阵
    SNR = SNR_i #200  # 信噪比

    # 互耦矩阵
    q = 5
    MCM = generate_mutual_coupling_matrix(M,q) #generate_complex_matrix(M) # 生成互耦矩阵
    #printD(MCM)
    c = MCM[:,0] 
    sig = MCM @ A @ S # 生成信号
    X = awgn(sig, SNR)  # 添加噪声后信号
    R = (X @ X.conj().T) / L  # 信号协方差矩阵
    D, En = eigh(R, subset_by_index=[0, M-K-1]) # 噪声子空间与噪声特征值
    search_grid = np.arange(-90, 90.1, 0.1)  # 搜索范围
    G,c1,tg_DOA,DOA_lst = huou_way(M,K,En,search_grid,MCM,c,q) # 互耦方法求DOA角度
    #doa_guji = DOA_lst[:] #find_top_n_angles(DOA_lst,K)
    doa_guji = DOA_lst[:K]
    doa_guji = sort_doa(doa_guji)
    real_doa = find_top_n_angles(theta,K)
    ## 计算 MCM
    mat_C = gMatC(M) # 初始化 MC 下标矩阵
    phi = gMat_J_all(M,mat_C,M-q)
    e = np.eye(M-q, 1)  # d 约束条件
    guji_mcm = Cal_mcm(phi,M,e,mat_C,En,doa_guji,M-q,K)
    #printD(guji_mcm[:,0])
    #printD(c)
    rmse_xiangliang = calculate_rmse_complex(guji_mcm[:,0],c)
    return rmse_xiangliang
    print(rmse_xiangliang)
    print(doa_guji)
    print(real_doa)
    rmse_i = calculate_rmse(doa_guji,real_doa)
    #print(rmse_i)
    #print(rmse_i)
   # return rmse_i
    # # 绘图
    G1 = 10 * np.log10(G / np.max(G))
    plt.figure()
    plt.title(str(SNR_i)+'dB '+str(q)+" "+str(M))
    plt.plot(search_grid, G1,label='g1')
    plt.legend()
    plt.grid(True)
    plt.show()
    
def compute_rmse(snr_i):
    rmse_trials = [main(snr_i) for _ in range(100)]
    avg_rmse_i = np.mean(rmse_trials)
    return avg_rmse_i

from concurrent.futures import ProcessPoolExecutor, as_completed   
if __name__=="__main__":
    #main(30)
    
    snr_values = np.arange(0, 20, 0.5)
    rmse_values = []
    max_workers = 12  # 最大工作进程数
    with ProcessPoolExecutor(max_workers=max_workers) as executor:
        results = list(executor.map(compute_rmse, snr_values))
    for snr_i, avg_rmse_i in zip(snr_values, results):
        rmse_values.append(avg_rmse_i)
        print(f"SNR: {snr_i}, RMSE: {avg_rmse_i}")

    plt.figure(figsize=(10, 6))
    plt.plot(snr_values, 10*np.log10(rmse_values), marker='o', linestyle='-', color='b')
    plt.xlabel('SNR (dB)')
    plt.ylabel('RMSE')
    plt.title('RMSE vs SNR')
    plt.grid(True)
    plt.show()
    
    # snr_values = range(10, 60, 1)  # SNR 从 10 到 59
    # rmse_values = []

    # # 计算不同 SNR 对应的 RMSE
    # for snr_i in snr_values:
    #     rmse_trials = [main(snr_i) for _ in range(10)]  # 进行 10 次计算
    #     avg_rmse_i = np.mean(rmse_trials)  # 取 10 次计算的平均值
    #     rmse_values.append(avg_rmse_i)
    #     print(f"SNR: {snr_i}, RMSE: {avg_rmse_i}")

    # # 使用 Matplotlib 绘制图像
    # plt.figure(figsize=(10, 6))
    # plt.plot(snr_values, 10*np.log10(rmse_values), marker='o', linestyle='-', color='b')
    # plt.xlabel('SNR (dB)')
    # plt.ylabel('RMSE')
    # plt.title('RMSE vs SNR')
    # plt.grid(True)
    # plt.show()
